"""
使用pytorch进行训练
f(x)=w1*x^2+w2*x+b
loss=(f(xn)-y)**2
"""
import random
import matplotlib.pyplot as plt
import torch

# 定义已知的数据集,training的dataset
x_data = [1.0, 2.0, 3.0]
y_data = [2.0, 4.0, 6.0]

# 定义权重w1
w1 = torch.Tensor([1])
# 声明需要计算梯度
w1.requires_grad = True
# 定义权重w2
w2 = torch.Tensor([1])
w2.requires_grad = True
# 定义权重b
b = torch.Tensor([1])
b.requires_grad = True


# 定义前馈
def forward(x):
    return w1 * x ** 2 + w2 * x + b


# 定义单样本的损失
def loss(x, y):
    y_hat = forward(x)
    return (y_hat - y) ** 2


# 记录损失
allLoss = []
# 记录训练轮数
allEpoch = []

bestVal = []
lastLoss = 100
count = 0
# 训练
for epoch in range(10000):
    # 使用随机单样本
        idx = random.randint(0, len(x_data) - 1)
        x = x_data[idx]
        y = y_data[idx]

    # 使用批量样本,可以采用随机小批量,这里因为数据少,所以使用所有的数据
    # for x, y in zip(x_data, y_data):
        # 记录epoch
        allEpoch.append(count)
        count += 1
        # 构建单样本损失图
        l = loss(x, y)
        # 记录损失
        allLoss.append(l.item())
        if l.item() < lastLoss:
            lastLoss = l.item()
            bestVal = [w1, w2, b]
        # 计算反向传播,获得梯度
        l.backward()
        # 更新权重
        w1.data -= 0.01 * w1.grad.data
        w2.data -= 0.01 * w2.grad.data
        b.data -= 0.01 * b.grad.data
        # 将梯度置零
        w1.grad.data.zero_()
        w2.grad.data.zero_()
        b.grad.data.zero_()
        # 打印
        print("Epoch: {}, Loss: {}".format(epoch, l.item()))

w1 = bestVal[0]
w2 = bestVal[1]
b = bestVal[2]

print(f"w1: {bestVal[0]},f w2: {bestVal[1]}f, b: {bestVal[2]}")
print("the value of 4 is:", forward(4).item())

# 绘制损失-轮数图
plt.plot(allEpoch, allLoss)
plt.ylabel("loss")
plt.xlabel("epoch")
plt.grid()
plt.show()
